# encoding:utf-8
import sys
import random
import math
import heapq
import os
from optparse import OptionParser
from custom_rand import CustomRand

class Flow:
	def __init__(self, src, dst, size, t):
		self.src, self.dst, self.size, self.t = src, dst, size, t
	def __str__(self):
		return "%d %d 3 100 %d %.9f"%(self.src, self.dst, self.size, self.t)

def translate_bandwidth(b):
	if b == None:
		return None
	if type(b)!=str:
		return None
	if b[-1] == 'G':
		return float(b[:-1])*1e9
	if b[-1] == 'M':
		return float(b[:-1])*1e6
	if b[-1] == 'K':
		return float(b[:-1])*1e3
	return float(b)

def poisson(lam):
	return -math.log(1-random.random())*lam

if __name__ == "__main__":
	port = 80
	parser = OptionParser()
	parser.add_option("-c", "--cdf", dest = "cdf_file", help = "the file name of the traffic size cdf", default = "uniform_distribution.txt")
	parser.add_option("-o", "--output", dest = "output", help = "the output file name", default = "tmp_traffic.txt")
	parser.add_option("-n", "--nhost", dest = "nhost", help = "number of hosts")
	parser.add_option("-l", "--load", dest = "load", help = "the percentage of the traffic load to the network capacity", default = "0")
	parser.add_option("-b", "--bandwidth", dest = "bandwidth", help = "the bandwidth of host link (G/M/K), by default 10G", default = "10G")
	parser.add_option("-t", "--time", dest = "time", help = "the total run time (s), by default 10", default = "10")
	parser.add_option("--incast_load", dest = "incast_load", help = "the percentage of incast traffic load to the network capacity, by default 0.02", default = "0")
	parser.add_option("--incast_size", dest = "incast_size", help = "the size of each incast flow in bytes, by default 500KB", default = "500000")
	parser.add_option("--incast_senders", dest = "incast_senders", help = "the number of senders in each incast", default = "0")
	parser.add_option("--wan_load", dest = "wan_load", help = "wan workload", default = "0")
	parser.add_option("--wan_bd", dest = "wan_bd", help = "wan bandwidth", default = "10G")
	parser.add_option("--cp", dest = "is_cp", help = "is copy reset to destination", default = 0)
	options,args = parser.parse_args()

	base_t = 2 * 1e9

	if not options.nhost:
		print "please use -n to enter number of hosts"
		sys.exit(0)
	
	nhost = int(options.nhost)
	incast_senders = int(options.incast_senders)
	load = float(options.load)
	incast_load = float(options.incast_load)
	wan_load = float(options.wan_load)
	fileName = os.path.join("../", "input", options.cdf_file)
	output = os.path.join("..", "output", options.output)
	bandwidth = int(translate_bandwidth(options.bandwidth))
	wan_bandwidth = int(translate_bandwidth(options.wan_bd))
	time = int(float(options.time)*1e9) # translates to ns	
	incast_size = int(options.incast_size)
	

	# Param Check
	if bandwidth == None:
		print("Error: bandwidth format incorrect")
		sys.exit(0)
	if incast_senders > nhost - 1:
		print("Error: Incast senders exceed the number of hosts")
		sys.exit(0)
	print("Number of hosts: %d"%nhost)
	print("DCN Load: %.2f%%"%(load * 100))
	print("DCN Bandwidth: %dGbps"%(bandwidth/1e9))
	if incast_load != 0:
		print("Incast Load: %.2f%%"%(incast_load * 100))
		print("Incast Size: %d"%incast_size)
		print("Incast Senders: %d"%incast_senders)
	if wan_load != 0:
		print("WAN Load: %.2f%%"%(wan_load * 100))
		print("WAN Bandwidth: %dGbps"%(wan_bandwidth/1e9))
	print("CDF file: %s"%fileName)
	print("Time: %dms"%(time/1e6))


	file = open(fileName,"r")
	lines = file.readlines()
	# read the cdf, save in cdf as [[x_i, cdf_i] ...]
	cdf = []
	for line in lines:
		x,y = map(float, line.strip().split(' '))
		cdf.append([x,y])
	wan_file = open("../input/Alibaba_WAN.txt", "r")
	wan_lines = wan_file.readlines()
	wan_cdf = []
	for line in wan_lines:
		x, y = map(float, line.strip().split(' '))
		wan_cdf.append([x, y])

	# create a custom random generator, which takes a cdf, and generate number according to the cdf
	customRand = CustomRand()
	if not customRand.setCdf(cdf):
		print "Error: Not valid DCN cdf"
		sys.exit(0)
	wanRand = CustomRand()
	if not wanRand.setCdf(wan_cdf):
		print "Error: Not valid WAN cdf"
		sys.exit(0)

	ofile = open(output, "w")
	flows = []
	dcn_flows = 0
	wan_flows = 0
	incast_flows = 0

	# generate DCN random flows（In DCN）
	if load != 0:
		avg = customRand.getAvg()
		avg_inter_arrival = 1000000000/(bandwidth*load/8./avg) # 1s/每秒生成几条流 = 每条流平均生成间隔时间
		print("\nDCN Average flow arrive: %.2fus"%(avg_inter_arrival/1000.))
		print("DCN Average flow size: %.2fKB"%(avg/1000.))

		host_list = [(base_t + int(poisson(avg_inter_arrival)), i) for i in range(nhost)]
		heapq.heapify(host_list)
		while len(host_list) > 0:
			t,src = host_list[0]
			inter_t = int(poisson(avg_inter_arrival))
			if src < nhost / 2:
				dst = random.randint(0, nhost/2-1)
				while (dst == src):
					dst = random.randint(0, nhost/2-1)
			else:
				dst = random.randint(nhost/2, nhost-1)
				while (dst == src):
					dst = random.randint(nhost/2, nhost-1)
			if (t + inter_t > time + base_t):
				heapq.heappop(host_list)
			else:
				size = int(customRand.rand())
				if size <= 0:
					size = 1
				dcn_flows += 1
				flows.append([src, dst, size, t * 1e-9])
				heapq.heapreplace(host_list, (t + inter_t, src))

	# generate incast flows（In DCN）
	if incast_load != 0:
		total_size = 0
		avg_incast_arrival = 1000000000/(bandwidth*nhost*incast_load/8./(incast_size*incast_senders))
		print("\nIncast Average flow arrive: %.2fus"%(avg_incast_arrival/1000.))
		t = base_t + int(poisson(avg_incast_arrival))
		while t < base_t + time:
			dst = random.randint(0, nhost-1)
			inter = int(poisson(avg_incast_arrival))
			if dst < nhost / 2:
				src = random.sample(range(nhost / 2), incast_senders)
				while dst in src:
					src.remove(dst)
					src.append(random.randint(0, (nhost / 2) - 1))
			else:
				src = random.sample(range(nhost / 2, nhost), incast_senders)
				while dst in src:
					src.remove(dst)
					src.append(random.randint(nhost / 2, nhost - 1))
			size = int(poisson(incast_size))
			for s in src:
				flows.append([s, dst, size, t * 1e-9])
				total_size += size
				incast_flows += 1
			t += inter

	# generate WAN flows（Between DCN）
	if wan_load != 0:
		wan_avg = wanRand.getAvg()
		avg_wan_arrival = 1000000000/(wan_bandwidth*wan_load/8./wan_avg)
		t = base_t + int(poisson(avg_wan_arrival))
		print("\nWAN Average flow arrive: %.2fus"%(avg_wan_arrival/1000.))
		print("WAN Average flow size: %.2fKB"%(wan_avg/1000.))
		while t < base_t + time:
			src = random.randint(0, nhost-1)
			if src < nhost / 2:
				dst = random.randint(nhost / 2, nhost - 1)
			else:
				dst = random.randint(0, nhost / 2 - 1)
			size = int(poisson(wan_avg))
			flows.append([src, dst, size, t * 1e-9])
			wan_flows += 1
			t += int(poisson(avg_wan_arrival))

	flows = sorted(flows, key=lambda x: x[3])
	n_flow = dcn_flows + incast_flows + wan_flows

	ofile.write("%d\n"%n_flow)
	for [src, dst, size, t] in flows:
		ofile.write("%d %d 3 100 %d %.9f\n"%(src, dst, size, t))
	ofile.close()

	print("\n--------------- Susseccfully generated traffic ---------------")
	print("Total number of flows: %d, dcn flows: %d, wan flows: %d, incast flows: %d"%(n_flow, dcn_flows, wan_flows, incast_flows))
	print("Output file: %s"%output)
	if incast_load != 0:
		print("Incast verify, total size: %d, network capacity: %f"%(total_size, bandwidth * nhost * incast_load * (time / 1e9) / 8))

	if bool(options.is_cp):
		os.system("cp %s ../../simulation/input/"%output)
		print("Copy to input folder")